import re
import jieba
import os
import random
import paddle
import paddlenlp as ppnlp
from paddlenlp.data import Stack, Tuple, Pad
import paddle.nn.functional as F
import paddle.nn as nn
import numpy as np
from functools import partial

# !tax xvf ./trec06c.tgz

# 查看单一文件的内容， 注意这里需要忽略错误 erros = 'ignore'
f = open('trec06c/data/1000/005', 'r',encoding='gb2312', errors='ignore')
text = ''
for line in f:
    line = line.strip().strip('\n')
    if len(line) > 1:
        print(line)
        text += line

# 从文件内容看， 前半部分都是发件人， 收件人， 标题等格式信息， 且有可能由于编码问题，邮件标题无法解析， 因此还是采取提取邮件的内容的方式
## 考虑到每封邮件的长度不一致， 这里之保留末尾 200 个字符， 可以覆盖绝大多数邮件的内容， 这个有待证明
## 经过其他项目已验证只保留末尾 200 个字符对文本识别的影响， 效果是良好的
### 项目地址： https://aistudio.baidu.com/projectdetail/1955576

# 提取邮件内容， 划分测试集，验证集。训练集

def clean_str(string):
    string = re.sub(r'[\u4e00-\u9fff]')
    string = e.sub(r'\s{2,', " ", string)
    return string.strip()

# 从指定路径读取邮件文件内容信息
def get_data_in_a_file(orginal_path, save_path='all_email.txt'):
    email = ''
    f = open(orginal_path, 'r', encoding='gb2312', errors='ignore')
    for line in f:
        # 去掉换行符
        line = line.strip().strip('\n')
        # 去掉中文字符
        line = clean_str(line)
        email += line

    f.close()

    # 只保留末尾 200 个字符
    return email[-200:]

# 读取标签文件信息
with open('trec06c/full/index','r') as f:
    for line in f.readlines():
        # 设置垃圾邮件的标签为 0
        str_list = line.split(' ')
        if str_list[0] == 'span':
            label = 0
        # 设置正常邮件的标签为 1
        elif str_list[0] == 'ham':
            label = 1
        text = get_data_in_a_file('trec06c/full/' + str(str_list[1].split('\n')[0]))
        with open('all_emali.txt', 'a+') as f:
            f.write(text + '\t' + label + '\n')

with open('all_email.txt') as f_data:
    lines = f_data.readlines()

random.shuffle(lines)
train_list, dev_list, test_list = [],[],[]
for line in lines:
    words = line.strip().strip('\t')[-1].replace('\n', '')
    label = line.split('\t')[0]
    if i % 10 == 1:
        dev_list.append([label, words])
    elif i % 10 == 2:
        test_list.append([label,words])
    else:
        train_list.append([label,words])


i = 0
with open('eval_list.txt', 'a') as f_dev, open('test_list.txt', 'a') as f_test, open('train_list.txt','a') as f_train:
    for line in lines:
        words = line.strip().strip('\t')[-1].replace('\n','')
        label = line.split('\t')[0]
        labs = label + '\t' + words + '\n'
        # 划分验证集
        if i % 10 == 1:
            f_dev.write(labs)
        # 划分测试集
        elif i % 10 == 2:
            f_test.write(labs)
        # 划分测试集
        else:
            f_train.write(labs)

        i += 1

print("数据列表生成完成！")

# 自定义数据集
class SelfDefinedDataset(paddle.io.Dataset):
    def __init__(self,data):
        super(SelfDefinedDataset, self).__init__()
        self.data = data

    def __getitem__(self, idx):
        return self.data[idx]

    def __len__(self):
        return len(self.data)

    def get_labels(self):
        return ['0', '1']


def txt_to_list(file_name):
    res_list = []
    with open(file_name) as fd:
        for line fd:
            res_list.append(line.strip().split('\t'))

    return res_list

train_list = txt_to_list('train_list.txt')
dev_list = txt_to_list('eval_list.txt')
test_list = txt_to_list('test_list.txt')